{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-11-10T13:08:37.314694Z",
     "start_time": "2025-11-10T13:08:23.642645Z"
    }
   },
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    " \n",
    "# 设置模型名称\n",
    "teacher_model_name = \"gpt2-medium\"\n",
    "student_model_name = \"gpt2\"  # 也可自定义更小的模型结构\n",
    "\n",
    "\n",
    "# 加载教师和学生模型\n",
    "teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)\n",
    "teacher_tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "teacher_model = GPT2LMHeadModel.from_pretrained(teacher_model_name)\n",
    "teacher_model.eval()  # 推理/评估模式\n",
    " \n",
    "student_tokenizer = AutoTokenizer.from_pretrained(student_model_name)\n",
    "student_tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "student_model = GPT2LMHeadModel.from_pretrained(student_model_name)\n",
    " \n",
    "# 简单数据集：使用 wikitext-2 做语言建模蒸馏演示\n",
    "dataset = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\", split='train')"
   ],
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-10T13:08:57.074755Z",
     "start_time": "2025-11-10T13:08:55.839701Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def tokenize_fn(examples):\n",
    "    return teacher_tokenizer(examples[\"text\"], truncation=True, max_length=128)\n",
    " \n",
    "tokenized_dataset = dataset.map(tokenize_fn, batched=True, remove_columns=[\"text\"])\n",
    " \n",
    "# PyTorch DataLoader\n",
    "def collate_fn(batch):\n",
    "    # 这里直接使用 teacher_tokenizer 的 pad 方法，也可用 student_tokenizer\n",
    "    return teacher_tokenizer.pad(batch, return_tensors=\"pt\")\n",
    " \n",
    "from torch.utils.data import DataLoader\n",
    " \n",
    "train_loader = DataLoader(tokenized_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)"
   ],
   "id": "636a74a69f23c2ba",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Map:   0%|          | 0/36718 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e3f2e412da6c42c29e44454fd90f77c0"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-10T13:09:01.102557Z",
     "start_time": "2025-11-10T13:09:01.084745Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch.nn.functional as F\n",
    " \n",
    "def distillation_loss_function(teacher_logits, student_logits, \n",
    "                               labels, \n",
    "                               alpha=0.5, temperature=2.0):\n",
    "    \"\"\"\n",
    "    teacher_logits, student_logits: (batch_size, seq_len, vocab_size)\n",
    "    labels: (batch_size, seq_len)\n",
    "    alpha: 权重，平衡真实任务损失 与 蒸馏损失\n",
    "    temperature: 蒸馏温度\n",
    "    \n",
    "    返回: total_loss\n",
    "    \"\"\"\n",
    "    # 1) LM 真实标签交叉熵\n",
    "    #    让学生在真实标签上也保持一定的准确度\n",
    "    #    -100 表示填充位置不计算loss\n",
    "    lm_loss = F.cross_entropy(\n",
    "        student_logits.view(-1, student_logits.size(-1)),\n",
    "        labels.view(-1),\n",
    "        ignore_index=-100\n",
    "    )\n",
    " \n",
    "    # 2) 蒸馏损失 (KL 散度)\n",
    "    #    对 teacher / student 的 logits 做 softmax with temperature\n",
    "    #    p(t) = softmax(teacher_logits / T)\n",
    "    #    q(s) = softmax(student_logits / T)\n",
    "    teacher_probs = F.log_softmax(teacher_logits / temperature, dim=-1)\n",
    "    student_probs = F.log_softmax(student_logits / temperature, dim=-1)\n",
    "    \n",
    "    distill_loss = F.kl_div(\n",
    "        student_probs, \n",
    "        teacher_probs.exp(),  # kl_div 需要 target 是概率分布 (非 log)\n",
    "        reduction='batchmean'\n",
    "    ) * (temperature**2)\n",
    " \n",
    "    total_loss = alpha * lm_loss + (1 - alpha) * distill_loss\n",
    "    return total_loss, lm_loss, distill_loss"
   ],
   "id": "e61acb2cf3c05bbe",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-10T13:09:03.803283Z",
     "start_time": "2025-11-10T13:09:02.369242Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch.optim as optim\n",
    " \n",
    "# 冻结教师模型，不参与训练\n",
    "for param in teacher_model.parameters():\n",
    "    param.requires_grad = False\n",
    " \n",
    "optimizer = optim.AdamW(student_model.parameters(), lr=1e-5)\n",
    " \n",
    "num_epochs = 1  # 简单跑1轮演示\n",
    "alpha = 0.5\n",
    "temperature = 2.0\n",
    " \n",
    "student_model.train()\n",
    " \n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_val = 0.0\n",
    "    for step, batch in enumerate(train_loader):\n",
    "        input_ids = batch[\"input_ids\"]\n",
    "        attention_mask = batch[\"attention_mask\"]\n",
    " \n",
    "        with torch.no_grad():\n",
    "            teacher_out = teacher_model(input_ids, attention_mask=attention_mask)\n",
    "            teacher_logits = teacher_out.logits\n",
    " \n",
    "        student_out = student_model(input_ids, attention_mask=attention_mask)\n",
    "        student_logits = student_out.logits\n",
    " \n",
    "        # labels 用来计算学生的 LM 任务损失\n",
    "        labels = input_ids.clone()\n",
    "        # 也可以把 padding位置设为 -100\n",
    "        labels[labels==teacher_tokenizer.pad_token_id] = -100\n",
    " \n",
    "        loss, lm_loss, distill_loss = distillation_loss_function(\n",
    "            teacher_logits, student_logits,\n",
    "            labels,\n",
    "            alpha=alpha, temperature=temperature\n",
    "        )\n",
    " \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    " \n",
    "        total_loss_val += loss.item()\n",
    " \n",
    "        if step % 100 == 0:\n",
    "            print(f\"Epoch {epoch}, Step {step}, Loss {loss.item():.4f}, LM {lm_loss.item():.4f}, KD {distill_loss.item():.4f}\")\n",
    " \n",
    "    avg_loss = total_loss_val / (step+1)\n",
    "    print(f\"Epoch {epoch} finished, avg loss = {avg_loss:.4f}\")"
   ],
   "id": "af8e37bd233f439c",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "index out of range in self",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mIndexError\u001B[0m                                Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[15], line 22\u001B[0m\n\u001B[0;32m     19\u001B[0m attention_mask \u001B[38;5;241m=\u001B[39m batch[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mattention_mask\u001B[39m\u001B[38;5;124m\"\u001B[39m]\n\u001B[0;32m     21\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m torch\u001B[38;5;241m.\u001B[39mno_grad():\n\u001B[1;32m---> 22\u001B[0m     teacher_out \u001B[38;5;241m=\u001B[39m \u001B[43mteacher_model\u001B[49m\u001B[43m(\u001B[49m\u001B[43minput_ids\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mattention_mask\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mattention_mask\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m     23\u001B[0m     teacher_logits \u001B[38;5;241m=\u001B[39m teacher_out\u001B[38;5;241m.\u001B[39mlogits\n\u001B[0;32m     25\u001B[0m student_out \u001B[38;5;241m=\u001B[39m student_model(input_ids, attention_mask\u001B[38;5;241m=\u001B[39mattention_mask)\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[0;32m   1497\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[0;32m   1498\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[0;32m   1499\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[0;32m   1500\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[1;32m-> 1501\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m forward_call(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n\u001B[0;32m   1502\u001B[0m \u001B[38;5;66;03m# Do not call functions when jit is used\u001B[39;00m\n\u001B[0;32m   1503\u001B[0m full_backward_hooks, non_full_backward_hooks \u001B[38;5;241m=\u001B[39m [], []\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\transformers\\models\\gpt2\\modeling_gpt2.py:1076\u001B[0m, in \u001B[0;36mGPT2LMHeadModel.forward\u001B[1;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001B[0m\n\u001B[0;32m   1068\u001B[0m \u001B[38;5;250m\u001B[39m\u001B[38;5;124mr\u001B[39m\u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[0;32m   1069\u001B[0m \u001B[38;5;124;03mlabels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\u001B[39;00m\n\u001B[0;32m   1070\u001B[0m \u001B[38;5;124;03m    Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set\u001B[39;00m\n\u001B[0;32m   1071\u001B[0m \u001B[38;5;124;03m    `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`\u001B[39;00m\n\u001B[0;32m   1072\u001B[0m \u001B[38;5;124;03m    are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`\u001B[39;00m\n\u001B[0;32m   1073\u001B[0m \u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[0;32m   1074\u001B[0m return_dict \u001B[38;5;241m=\u001B[39m return_dict \u001B[38;5;28;01mif\u001B[39;00m return_dict \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;28;01melse\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mconfig\u001B[38;5;241m.\u001B[39muse_return_dict\n\u001B[1;32m-> 1076\u001B[0m transformer_outputs \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mtransformer\u001B[49m\u001B[43m(\u001B[49m\n\u001B[0;32m   1077\u001B[0m \u001B[43m    \u001B[49m\u001B[43minput_ids\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1078\u001B[0m \u001B[43m    \u001B[49m\u001B[43mpast_key_values\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mpast_key_values\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1079\u001B[0m \u001B[43m    \u001B[49m\u001B[43mattention_mask\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mattention_mask\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1080\u001B[0m \u001B[43m    \u001B[49m\u001B[43mtoken_type_ids\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mtoken_type_ids\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1081\u001B[0m \u001B[43m    \u001B[49m\u001B[43mposition_ids\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mposition_ids\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1082\u001B[0m \u001B[43m    \u001B[49m\u001B[43mhead_mask\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mhead_mask\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1083\u001B[0m \u001B[43m    \u001B[49m\u001B[43minputs_embeds\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs_embeds\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1084\u001B[0m \u001B[43m    \u001B[49m\u001B[43mencoder_hidden_states\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mencoder_hidden_states\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1085\u001B[0m \u001B[43m    \u001B[49m\u001B[43mencoder_attention_mask\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mencoder_attention_mask\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1086\u001B[0m \u001B[43m    \u001B[49m\u001B[43muse_cache\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43muse_cache\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1087\u001B[0m \u001B[43m    \u001B[49m\u001B[43moutput_attentions\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43moutput_attentions\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1088\u001B[0m \u001B[43m    \u001B[49m\u001B[43moutput_hidden_states\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43moutput_hidden_states\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1089\u001B[0m \u001B[43m    \u001B[49m\u001B[43mreturn_dict\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mreturn_dict\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1090\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m   1091\u001B[0m hidden_states \u001B[38;5;241m=\u001B[39m transformer_outputs[\u001B[38;5;241m0\u001B[39m]\n\u001B[0;32m   1093\u001B[0m \u001B[38;5;66;03m# Set device for model parallelism\u001B[39;00m\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[0;32m   1497\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[0;32m   1498\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[0;32m   1499\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[0;32m   1500\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[1;32m-> 1501\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m forward_call(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n\u001B[0;32m   1502\u001B[0m \u001B[38;5;66;03m# Do not call functions when jit is used\u001B[39;00m\n\u001B[0;32m   1503\u001B[0m full_backward_hooks, non_full_backward_hooks \u001B[38;5;241m=\u001B[39m [], []\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\transformers\\models\\gpt2\\modeling_gpt2.py:843\u001B[0m, in \u001B[0;36mGPT2Model.forward\u001B[1;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict)\u001B[0m\n\u001B[0;32m    840\u001B[0m head_mask \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mget_head_mask(head_mask, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mconfig\u001B[38;5;241m.\u001B[39mn_layer)\n\u001B[0;32m    842\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m inputs_embeds \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[1;32m--> 843\u001B[0m     inputs_embeds \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mwte\u001B[49m\u001B[43m(\u001B[49m\u001B[43minput_ids\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    844\u001B[0m position_embeds \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mwpe(position_ids)\n\u001B[0;32m    845\u001B[0m hidden_states \u001B[38;5;241m=\u001B[39m inputs_embeds \u001B[38;5;241m+\u001B[39m position_embeds\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[0;32m   1497\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[0;32m   1498\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[0;32m   1499\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[0;32m   1500\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[1;32m-> 1501\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m forward_call(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n\u001B[0;32m   1502\u001B[0m \u001B[38;5;66;03m# Do not call functions when jit is used\u001B[39;00m\n\u001B[0;32m   1503\u001B[0m full_backward_hooks, non_full_backward_hooks \u001B[38;5;241m=\u001B[39m [], []\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\torch\\nn\\modules\\sparse.py:162\u001B[0m, in \u001B[0;36mEmbedding.forward\u001B[1;34m(self, input)\u001B[0m\n\u001B[0;32m    161\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;21mforward\u001B[39m(\u001B[38;5;28mself\u001B[39m, \u001B[38;5;28minput\u001B[39m: Tensor) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m Tensor:\n\u001B[1;32m--> 162\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mF\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43membedding\u001B[49m\u001B[43m(\u001B[49m\n\u001B[0;32m    163\u001B[0m \u001B[43m        \u001B[49m\u001B[38;5;28;43minput\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mweight\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mpadding_idx\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mmax_norm\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    164\u001B[0m \u001B[43m        \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mnorm_type\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mscale_grad_by_freq\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43msparse\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[1;32mE:\\ProgramData\\anaconda3\\envs\\chatglm3-qlora\\lib\\site-packages\\torch\\nn\\functional.py:2210\u001B[0m, in \u001B[0;36membedding\u001B[1;34m(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)\u001B[0m\n\u001B[0;32m   2204\u001B[0m     \u001B[38;5;66;03m# Note [embedding_renorm set_grad_enabled]\u001B[39;00m\n\u001B[0;32m   2205\u001B[0m     \u001B[38;5;66;03m# XXX: equivalent to\u001B[39;00m\n\u001B[0;32m   2206\u001B[0m     \u001B[38;5;66;03m# with torch.no_grad():\u001B[39;00m\n\u001B[0;32m   2207\u001B[0m     \u001B[38;5;66;03m#   torch.embedding_renorm_\u001B[39;00m\n\u001B[0;32m   2208\u001B[0m     \u001B[38;5;66;03m# remove once script supports set_grad_enabled\u001B[39;00m\n\u001B[0;32m   2209\u001B[0m     _no_grad_embedding_renorm_(weight, \u001B[38;5;28minput\u001B[39m, max_norm, norm_type)\n\u001B[1;32m-> 2210\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mtorch\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43membedding\u001B[49m\u001B[43m(\u001B[49m\u001B[43mweight\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43minput\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mpadding_idx\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mscale_grad_by_freq\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43msparse\u001B[49m\u001B[43m)\u001B[49m\n",
      "\u001B[1;31mIndexError\u001B[0m: index out of range in self"
     ]
    }
   ],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-10T13:13:57.130443Z",
     "start_time": "2025-11-10T13:13:57.071260Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 确保pad_token_id有效\n",
    "pad_token_id = teacher_tokenizer.pad_token_id if teacher_tokenizer.pad_token_id is not None else teacher_tokenizer.eos_token_id\n",
    "\n",
    "# 安全设置padding标签\n",
    "labels = input_ids.clone()\n",
    "labels[labels == pad_token_id] = -100"
   ],
   "id": "f2e6b6201ab70770",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-10T13:14:13.683229Z",
     "start_time": "2025-11-10T13:14:13.676219Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "def distillation_loss_function(teacher_logits, student_logits, labels, alpha=0.5, temperature=2.0):\n",
    "    # 知识蒸馏损失\n",
    "    distill_loss = nn.KLDivLoss()(\n",
    "        torch.nn.functional.log_softmax(student_logits/temperature, dim=-1),\n",
    "        torch.nn.functional.softmax(teacher_logits/temperature, dim=-1)\n",
    "    ) * (temperature**2)\n",
    "    \n",
    "    # 学生模型的交叉熵损失\n",
    "    ce_loss = nn.CrossEntropyLoss(ignore_index=-100)(student_logits.view(-1, student_logits.size(-1)), labels.view(-1))\n",
    "    \n",
    "    # 组合损失\n",
    "    total_loss = alpha * ce_loss + (1 - alpha) * distill_loss\n",
    "    return total_loss, ce_loss, distill_loss"
   ],
   "id": "454b18b7f9f2c2e2",
   "outputs": [],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-10T13:14:47.331386Z",
     "start_time": "2025-11-10T13:14:47.301860Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 标签处理修正\n",
    "pad_token_id = teacher_tokenizer.pad_token_id or teacher_tokenizer.eos_token_id\n",
    "labels = input_ids.clone()\n",
    "labels[labels == pad_token_id] = -100\n",
    "\n",
    "# 使用标准蒸馏损失函数\n",
    "loss, lm_loss, distill_loss = distillation_loss_function(\n",
    "    teacher_logits, \n",
    "    student_logits,\n",
    "    labels,\n",
    "    alpha=alpha, \n",
    "    temperature=temperature\n",
    ")"
   ],
   "id": "27a0aecd2ba9933",
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'teacher_logits' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[18], line 8\u001B[0m\n\u001B[0;32m      4\u001B[0m labels[labels \u001B[38;5;241m==\u001B[39m pad_token_id] \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m100\u001B[39m\n\u001B[0;32m      6\u001B[0m \u001B[38;5;66;03m# 使用标准蒸馏损失函数\u001B[39;00m\n\u001B[0;32m      7\u001B[0m loss, lm_loss, distill_loss \u001B[38;5;241m=\u001B[39m distillation_loss_function(\n\u001B[1;32m----> 8\u001B[0m     \u001B[43mteacher_logits\u001B[49m, \n\u001B[0;32m      9\u001B[0m     student_logits,\n\u001B[0;32m     10\u001B[0m     labels,\n\u001B[0;32m     11\u001B[0m     alpha\u001B[38;5;241m=\u001B[39malpha, \n\u001B[0;32m     12\u001B[0m     temperature\u001B[38;5;241m=\u001B[39mtemperature\n\u001B[0;32m     13\u001B[0m )\n",
      "\u001B[1;31mNameError\u001B[0m: name 'teacher_logits' is not defined"
     ]
    }
   ],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "6713a14114561bcd"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
